package com.google.code.coma.loadbalancer.space;

import com.google.code.coma.Coma;
import com.google.code.coma.domain.Part;
import com.google.code.coma.loadbalancer.LoadBalancer;
import com.google.code.coma.loadbalancer.entity.Transfer;
import com.google.code.coma.space.IdSelectorCallback;
import com.google.code.coma.space.SelectAllSelectorCallback;
import com.google.code.coma.space.SpaceSelector;
import com.google.code.coma.space.TypeAmountPartSelectorCallback;
import com.google.code.coma.space.util.SpaceManager;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.mozartspaces.core.*;

import java.net.URI;
import java.util.*;

/**
 * @author Dominik Strasser, dominikstr@gmail.com
 */
public class SpaceLoadBalancer extends LoadBalancer {
	private static final Log log = LogFactory.getLog(SpaceLoadBalancer.class);

	private List<String> factories;

	private URI spaceURI;

	private MzsCore core;

	private Capi capi;

	public SpaceLoadBalancer(String spaceURI, String... factories) {
		this.factories = Arrays.asList(factories);
		this.spaceURI = URI.create(spaceURI);
		core = DefaultMzsCore.newInstance(0);
		capi = new Capi(core);
	}

	@Override
	public void run() {
		Map<String, List<Part>> factoriesParts = new HashMap<String, List<Part>>();
		TransactionReference tx;
		try {
			tx = capi.createTransaction(MzsConstants.TransactionTimeout.INFINITE, spaceURI);
		} catch (MzsCoreException e) {
			log.error("Unable to create transaction", e);
			return;
		}
		for (String factoriy : factories) {
			try {
				ContainerReference containerReference = SpaceManager.getOrCreateNamedContainer(spaceURI, factoriy + "_" + Coma.PARTS_CONTAINER, capi);
				factoriesParts.put(factoriy,
						capi.<Part>read(containerReference, new SpaceSelector(new SelectAllSelectorCallback()), MzsConstants.RequestTimeout.DEFAULT, tx));
			} catch (MzsCoreException e) {
				try {
					capi.rollbackTransaction(tx);
				} catch (MzsCoreException e1) {
					log.error("Error in transaction rollback", e1);
				}
				throw new RuntimeException(e);
			}
		}

		Transfer transfer = getLoadBalanceTransfers(factoriesParts);

		if (transfer != null) {
			try {
				ContainerReference from = SpaceManager.getOrCreateNamedContainer(spaceURI, transfer.getSupply().getFactoryId() + "_" + Coma.PARTS_CONTAINER, capi);
				ContainerReference to = SpaceManager.getOrCreateNamedContainer(spaceURI, transfer.getDemand().getFactoryId() + "_" + Coma.PARTS_CONTAINER, capi);

                int amount = transfer.getSupply().getAmount() > transfer.getDemand().getAmount() ? transfer.getDemand().getAmount() : transfer.getSupply().getAmount();

				List<Part> parts = capi.take(from, new SpaceSelector(new TypeAmountPartSelectorCallback(amount, transfer.getSupply().getType())), MzsConstants.RequestTimeout.INFINITE, tx);
				for (Part p : parts) {
					capi.write(new Entry(p), to, MzsConstants.RequestTimeout.DEFAULT, tx);
				}
				if (log.isInfoEnabled()) {
					log.info("Moving " + parts.size() + " " + transfer.getDemand().getType() + "s from " + transfer.getSupply().getFactoryId() + " to " + transfer.getDemand().getFactoryId());
				}
			} catch (MzsCoreException e) {
				log.error("Error fetching container", e);
				return;
			}
		}

		try {
			capi.commitTransaction(tx);
		} catch (MzsCoreException e1) {
			log.error("Error in transaction commit", e1);
		}
	}
}
